% 清理工作区和命令窗口，关闭所有图形窗口
clearvars;
close all;
clc;

% --- 1. 自动查找并加载预测数据 ---
predictions_folder = 'LSTM_predictions';
fprintf('正在 "%s" 文件夹中查找 .mat 文件...\n', predictions_folder);

% 检查预测文件夹是否存在
if ~isfolder(predictions_folder)
    error('错误: 找不到名为 "%s" 的文件夹。\n请确保此脚本与 "%s" 文件夹在同一目录下。', predictions_folder, predictions_folder);
end

% 在文件夹中查找所有 .mat 文件
mat_files = dir(fullfile(predictions_folder, '*.mat'));

% 检查是否找到了任何 .mat 文件
if isempty(mat_files)
    error('错误: 在 "%s" 文件夹中没有找到任何 .mat 文件。', predictions_folder);
end

% 加载找到的第一个 .mat 文件
mat_file_to_load = fullfile(predictions_folder, mat_files(1).name);
fprintf('正在加载数据文件: %s\n', mat_files(1).name);
data = load(mat_file_to_load);
fprintf('数据加载成功。\n');

% 从加载的结构体中提取变量，方便后续使用
prediction_base_time = data.prediction_base_time;
drone_id = data.drone_id;
predicted_time = data.predicted_time;
pred_x = data.pred_x;
pred_y = data.pred_y;
pred_z = data.pred_z;


% --- 2. 数据预处理 ---
% 获取所有唯一的预测起始时间和无人机ID
% unique() 函数会自动对结果进行排序
unique_base_times = unique(prediction_base_time);
unique_drone_ids = unique(drone_id);

num_frames = length(unique_base_times);
num_drones = length(unique_drone_ids);

fprintf('共找到 %d 个动画帧 (不同的预测起始时间)。\n', num_frames);
fprintf('共找到 %d 架无人机。\n', num_drones);


% --- 3. 设置 3D 可视化窗口 ---
fprintf('正在创建可视化窗口...\n');
fig = figure('Name', '无人机轨迹预测可视化 (MATLAB)', 'Position', [100, 100, 1000, 800]);
ax = axes('Parent', fig);
hold(ax, 'on');
grid(ax, 'on');
axis(ax, 'equal'); % 使各轴比例相同，避免轨迹变形
view(3); % 设置为 3D 视角
rotate3d on; % 允许用鼠标旋转视角

xlabel(ax, 'X 坐标 (米)', 'FontSize', 12);
ylabel(ax, 'Y 坐标 (米)', 'FontSize', 12);
zlabel(ax, 'Z 坐标 (米)', 'FontSize', 12);

% 为不同的无人机设置不同的颜色
colors = lines(num_drones);

% 初始化绘图句柄，这是实现高效动画的关键
plot_handles = gobjects(num_drones, 1);
legend_entries = cell(num_drones, 1);

for i = 1:num_drones
    % 初始化一个空的3D线条对象
    plot_handles(i) = plot3(ax, NaN, NaN, NaN, ...
        'LineWidth', 2, ...
        'Color', colors(i,:), ...
        'Marker', '.', ...
        'MarkerSize', 10);
    % MATLAB 中习惯从1开始编号，所以将ID+1显示
    legend_entries{i} = sprintf('无人机 %d', unique_drone_ids(i) + 1);
end

legend(ax, plot_handles, legend_entries, 'FontSize', 12, 'Location', 'northeast');


% --- 4. 循环播放动画 ---
fprintf('开始播放预测动画...\n');

for i = 1:num_frames
    current_base_time = unique_base_times(i);
    
    % 更新图标题，显示当前是哪个时间点的预测
    title(ax, sprintf('显示在 %.2fs 做出的轨迹预测', current_base_time), 'FontSize', 16);
    
    % 为每架无人机更新其预测轨迹
    for j = 1:num_drones
        current_drone_id = unique_drone_ids(j);
        
        % 找到当前帧和当前无人机对应的所有数据点的索引
        idx = (prediction_base_time == current_base_time) & (drone_id == current_drone_id);
        
        % 提取对应的坐标
        x_coords = pred_x(idx);
        y_coords = pred_y(idx);
        z_coords = pred_z(idx);
        
        % 更新绘图句柄的数据，而不是重新绘图
        set(plot_handles(j), 'XData', x_coords, 'YData', y_coords, 'ZData', z_coords);
    end
    
    % 强制 MATLAB 立即绘制更新后的图形
    drawnow;
    
    % 暂停一段时间，控制动画速度
    pause(0.1); 
end

fprintf('动画播放完成。\n');
